package com.adv.ems.db.utils;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.ArrayUtil;
import lombok.Data;
import java.sql.*;
import java.util.*;
import java.util.Date;
import java.util.stream.Collectors;

public class DbUtil {
    private static Map<String, Class> dbTypeMapJavaType = new HashMap<>();
    static {
        dbTypeMapJavaType.put("bigint", Long.class);
        dbTypeMapJavaType.put("char", Character.class);
        dbTypeMapJavaType.put("date", Date.class);
        dbTypeMapJavaType.put("datetime", Date.class);
        dbTypeMapJavaType.put("decimal", Double.class);
        dbTypeMapJavaType.put("double", Double.class);
        dbTypeMapJavaType.put("enum", String.class);
        dbTypeMapJavaType.put("float", Float.class);
        dbTypeMapJavaType.put("int", Integer.class);
        dbTypeMapJavaType.put("json", String.class);
        dbTypeMapJavaType.put("longtext", String.class);
        dbTypeMapJavaType.put("mediumtext", String.class);
        dbTypeMapJavaType.put("text", String.class);
        dbTypeMapJavaType.put("smallint", Integer.class);
        dbTypeMapJavaType.put("timestamp", Date.class);
        dbTypeMapJavaType.put("tinyint", Integer.class);
        dbTypeMapJavaType.put("varchar", String.class);
        dbTypeMapJavaType.put("longblob", Byte.class);
        try {
            Class.forName("com.mysql.cj.jdbc.Driver");
        } catch (ClassNotFoundException e) {
            throw new RuntimeException(e);
        }
    }

    public static Map<String, List<ColumnInfo>> readInfo(String ip, String port, String username, String pwd, String dbName, String...locateTables) {
        Connection connection = null;
        PreparedStatement preparedStatement = null;
        ResultSet resultSet = null;
        Map<String, List<ColumnInfo>> tableInfo = new HashMap<String, List<ColumnInfo>>();
        try {
            String url = String.format("jdbc:mysql://%s:%s/information_schema?useUnicode=true&characterEncoding=utf8", ip, port);
            connection = DriverManager.getConnection(url, username, pwd);
            StringBuilder sqlBuilder = new StringBuilder();
            sqlBuilder.append("SELECT table_name,column_name,data_type,column_key FROM COLUMNS WHERE ");
            sqlBuilder.append(String.format(" TABLE_SCHEMA = '%s' ", dbName));
            if (ArrayUtil.isNotEmpty(locateTables)) {
                sqlBuilder.append(" AND TABLE_NAME IN ('");
                sqlBuilder.append(CollectionUtil.join(Arrays.stream(locateTables).collect(Collectors.toList()), "','"));
                sqlBuilder.append("')");
            }
            sqlBuilder.append(" ORDER BY table_name ASC,ORDINAL_POSITION ASC ");
            preparedStatement = connection.prepareStatement(sqlBuilder.toString());
            resultSet = preparedStatement.executeQuery();
            while (resultSet.next()) {
                String tableName = resultSet.getString("table_name");
                List<ColumnInfo> columnInfos = tableInfo.getOrDefault(tableName, new ArrayList<>());
                String columnName = resultSet.getString("column_name");
                String dataType = resultSet.getString("data_type");
                String columnKey = resultSet.getString("column_key");
                ColumnInfo  columnInfo = new ColumnInfo();
                columnInfo.setColumnName(columnName);
                columnInfo.setColumnType(dataType);
                columnInfo.setIsPrimaryKey(columnKey != null && columnKey.equals("PRI"));
                if (dbTypeMapJavaType.get(dataType) == null) {
                    throw new RuntimeException(String.format("Unknown datatype '%s'", dataType));
                }
                columnInfo.setJavaType(dbTypeMapJavaType.get(dataType));
                columnInfos.add(columnInfo);
                tableInfo.put(tableName, columnInfos);
            }
            return tableInfo;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }finally {
            try {
                if (resultSet != null) {
                    resultSet.close();
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            try {
                if (preparedStatement != null) {
                    preparedStatement.close();
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            try {
                if (connection != null) {
                    connection.close();
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
    }

    @Data
    public static class ColumnInfo {
        private String columnName;
        private String columnType;
        private Class javaType;
        private Boolean isPrimaryKey;
    }
}
